"""
预测入口 - 对新数据进行预测
"""
import argparse
import torch
import numpy as np
import pandas as pd
from pathlib import Path
from tqdm import tqdm

from configs.default import DefaultConfig
from models import build_model
from data.dataset import MeniscusDataset
from data.utils import load_csv_data
from utils.checkpoint import load_checkpoint
from utils.general import set_seed, colorstr, select_device


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Predict with Meniscus Volume Model')
    
    parser.add_argument('--weights', type=str, required=True,
                       help='path to model weights')
    parser.add_argument('--source', type=str, required=True,
                       help='source csv file or directory')
    parser.add_argument('--output', type=str, default=None,
                       help='output file path (default: source_predictions.csv)')
    parser.add_argument('--batch-size', type=int, default=32,
                       help='batch size')
    parser.add_argument('--device', type=str, default='',
                       help='cuda device')
    parser.add_argument('--save-results', action='store_true',
                       help='save prediction results to csv')
    
    return parser.parse_args()


def predict(model, data, config, device):
    """
    对数据进行预测
    
    参数:
        model: 模型
        data: 数据列表
        config: 配置
        device: 设备
    
    返回:
        predictions: 预测结果列表
        targets: 真实标签列表
    """
    from torch.utils.data import DataLoader
    
    # 创建数据集
    dataset = MeniscusDataset(
        data,
        image_size=config.image_size,
        augment=False,
        filter_negative=False
    )
    
    # 创建数据加载器
    dataloader = DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=False,
        num_workers=config.num_workers,
        pin_memory=True if torch.cuda.is_available() else False
    )
    
    # 预测
    model.eval()
    predictions = []
    targets = []
    
    print(colorstr('bright_green', '\nPredicting...'))
    pbar = tqdm(dataloader, desc='Prediction', ncols=100)
    
    with torch.no_grad():
        for images, scales, features, target in pbar:
            images = images.to(device)
            scales = scales.to(device)
            features = features.to(device)
            
            # 预测
            output = model(images, scales, features)
            
            predictions.extend(output.cpu().numpy().flatten())
            targets.extend(target.numpy().flatten())
    
    return predictions, targets


def save_predictions(predictions, targets, output_path):
    """
    保存预测结果
    
    参数:
        predictions: 预测值列表
        targets: 真实值列表
        output_path: 输出路径
    """
    df = pd.DataFrame({
        'True_Volume': targets,
        'Predicted_Volume': predictions,
        'Error': np.array(predictions) - np.array(targets),
        'Absolute_Error': np.abs(np.array(predictions) - np.array(targets)),
        'Relative_Error_%': np.abs((np.array(predictions) - np.array(targets)) / (np.array(targets) + 1e-8)) * 100
    })
    
    df.to_csv(output_path, index=False, encoding='utf-8-sig')
    print(colorstr('bright_green', f'\n✓ Results saved to: {output_path}'))
    
    # 打印统计信息
    print(f"\nPrediction Statistics:")
    print(f"  Samples: {len(predictions)}")
    print(f"  Mean Error: {df['Error'].mean():.4f}")
    print(f"  MAE: {df['Absolute_Error'].mean():.4f}")
    print(f"  MAPE: {df['Relative_Error_%'].mean():.2f}%")


def main():
    """主函数"""
    # 解析参数
    args = parse_args()
    
    # 检查权重文件
    weights_path = Path(args.weights)
    if not weights_path.exists():
        print(colorstr('red', f'Error: Weights file not found: {weights_path}'))
        return
    
    # 检查源文件
    source_path = Path(args.source)
    if not source_path.exists():
        print(colorstr('red', f'Error: Source file not found: {source_path}'))
        return
    
    print(colorstr('bright_blue', 'bold', '\n' + '='*70))
    print(colorstr('bright_blue', 'bold', 'Model Prediction'.center(70)))
    print(colorstr('bright_blue', 'bold', '='*70))
    
    # 加载checkpoint
    checkpoint = torch.load(weights_path, map_location='cpu')
    
    # 创建配置
    config = DefaultConfig()
    if 'config' in checkpoint:
        for k, v in checkpoint['config'].items():
            if hasattr(config, k):
                setattr(config, k, v)
    
    config.batch_size = args.batch_size
    config.device = select_device(args.device)
    
    # 打印信息
    print(f"\nWeights: {weights_path}")
    print(f"Source: {source_path}")
    print(f"Model Type: {config.model_type}")
    print(f"Device: {config.device}")
    
    # 构建模型
    print(colorstr('bright_green', '\nBuilding model...'))
    model = build_model(config).to(config.device)
    
    # 加载权重
    load_checkpoint(weights_path, model, device=config.device)
    
    # 加载数据
    print(colorstr('bright_green', '\nLoading data...'))
    if source_path.is_file():
        data = load_csv_data(source_path.parent, [source_path.name])
    else:
        csv_files = list(source_path.glob('*.csv'))
        data = load_csv_data(source_path, [f.name for f in csv_files])
    
    print(f"Loaded {len(data)} samples")
    
    # 预测
    predictions, targets = predict(model, data, config, config.device)
    
    # 保存结果
    if args.save_results or args.output:
        output_path = args.output if args.output else source_path.parent / f'{source_path.stem}_predictions.csv'
        save_predictions(predictions, targets, output_path)
    
    print(colorstr('bright_green', 'bold', '\n✓ Prediction completed!'))


if __name__ == '__main__':
    main()